
import math


w_i = 0  # 坡道附加阻力


class Locomotive:
    def __init__(self) -> None:
        self.power = 9600  # 功率，单位为千瓦
        self.startSpeed = 20  # 具备启动牵引力的最大速度
        self.adhesionMass = 150  # 黏着重量，单位为吨
        self.gravity = 9.81  # 重力加速度，单位为m/s^2

    def adhesion_coefficient(self, v: float) -> float:
        if v < self.startSpeed:
            v = self.startSpeed
        return 0.4082 - 0.0021 * v

    def max_traction(self, v: float) -> float:
        """
        计算机车的最大牵引力

        机车在坡道上的黏着系数会减少，因此最大牵引力也会减小。
        """
        mu = self.adhesion_coefficient(v)  # 计算当前速度下的黏着系数

        # 计算坡道对牵引力造成的影响
        slope_coe = math.cos(math.atan(0.001*w_i))

        # 计算黏着牵引力限制，单位kN。t和kN的乘1000单位转换刚好约分
        max_adhesion_traction = mu * self.adhesionMass * self.gravity*slope_coe

        # 如果速度小于启动速度，那么就直接返回。
        # 这一步最主要的目的是避免除零
        if v <= self.startSpeed:
            return max_adhesion_traction

        # 计算功率的限制，单位kN。F=P/V，其中kW和kN的1000刚好约分，但需要将速度转换为m/s。
        max_power_traction = self.power / (v / 3.6)
        return min(max_adhesion_traction, max_power_traction)

    def friction(self, v: float) -> float:
        """
        机车的阻力,返回单位为kN
        """
        return (2.40+0.0022*v+0.000391*v**2+w_i)*(self.adhesionMass*self.gravity)/1000
